from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# 模型配置
MODEL_PATH = "./llama.cpp/deepseek-r1-1.5b.q4_0.gguf"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# 加载模型和分词器
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_PATH,
    torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32,
    device_map=DEVICE
)

# 定义请求体
class Request(BaseModel):
    prompt: str
    max_length: int = 512
    temperature: float = 0.7

# 初始化FastAPI
app = FastAPI()

# 允许跨域
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

@app.post("/generate")
async def generate_text(request: Request):
    try:
        # 编码输入
        inputs = tokenizer(
            request.prompt,
            return_tensors="pt"
        ).to(DEVICE)

        # 生成文本
        outputs = model.generate(
            **inputs,
            max_length=request.max_length,
            temperature=request.temperature,
            top_p=0.9,
            repetition_penalty=1.2,
            do_sample=True
        )

        # 解码输出
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        return {"result": response}

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)
